/*
 Xern Entity Relocation Network
 Copyright (C) 2011 Ronie Salgado <roniesalg@gmail.com>

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
*/

#include <cstdio>
#include <boost/uuid/uuid_io.hpp>
#include <boost/lexical_cast.hpp>
#include "Xern/IDataFilter.hpp"
#include "Xern/IResourceCache.hpp"
#include "Xern/DataFilter.hpp"
#include "Xern/FileSystem.hpp"
#include "Xern/ResourceManager.hpp"
#include "Xern/Engine.hpp"
#include "Xern/Exceptions.hpp"
#include "Xern/Log.hpp"
#include "TransferManager.hpp"

namespace XernProtocol
{
	// TransferState
	class TransferState
	{
	public:
		enum State
		{
			ST_WAITING,
			ST_RECEIVING,
		};

		TransferState();
		~TransferState();

		typedef std::set<boost::shared_ptr<IResourceLoadedCallback> > Callbacks;
		typedef std::vector<boost::shared_ptr<ResourceData> > EarlyData;

		State state;
		Callbacks callbacks;

		boost::shared_ptr<IDataFilter> stream;
		std::string path;
		std::string mime;
		Sha256 hash;
		int numblocks;
		int remainingBlocks;
		size_t fileSize;
		size_t blockSize;
	};

	TransferState::TransferState()
	{
		state = ST_WAITING;
	}

	TransferState::~TransferState()
	{
	}

	// TransferManager
	TransferManager::TransferManager(boost::shared_ptr<Session> session)
		: session(session)
	{
	}

	TransferManager::~TransferManager()
	{
	}

	bool TransferManager::ProcessMessage(boost::shared_ptr<Message> message)
	{
		switch(message->GetType())
		{
		case MT_RESOURCE_REQUEST:
			OnResourceRequest(boost::static_pointer_cast<ResourceRequest> (message));
			return true;
		case MT_RESOURCE_REQUEST_DATA:
			OnResourceRequestData(boost::static_pointer_cast<ResourceRequestData> (message));
			return true;
		case MT_RESOURCE_DATA:
			OnResourceData(boost::static_pointer_cast<ResourceData> (message));
			return true;
		case MT_RESOURCE_DENIAL:
			OnResourceDenial(boost::static_pointer_cast<ResourceDenial> (message));
			return true;
		case MT_RESOURCE_DESCRIPTION:
			OnResourceDescription(boost::static_pointer_cast<ResourceDescription> (message));
			return true;
		default:
			return false;
		}
	}

	void TransferManager::DownloadResource(const uuid &id, boost::shared_ptr<IResourceLoadedCallback> callback)
	{
		// Check for an existing download.
		TransferStates::iterator it = transferStates.find(id);
		if(it != transferStates.end())
		{
			// Add the callback to the existing download.
			boost::shared_ptr<TransferState> state = it->second;
			state->callbacks.insert(callback);
			return;
		}

		// Create a new download.
		boost::shared_ptr<TransferState> state(new TransferState());
		state->state = TransferState::ST_WAITING;
		state->callbacks.insert(callback);
		transferStates.insert(std::make_pair(id, state));

		// Send the request message.
		boost::shared_ptr<ResourceRequest> request(new ResourceRequest());
		request->resourceId = id;
		session->Send(request);
	}

	void TransferManager::OnResourceRequest(boost::shared_ptr<ResourceRequest> message)
	{
		// Find the resource.
		boost::shared_ptr<ResourceManager> resourceManager = Engine::Get().GetResourceManager();
		boost::shared_ptr<IResource> resource = resourceManager->FindResource(message->resourceId);
		if(!resource)
		{
			// Couldn't find the resource, send a denial message.
			boost::shared_ptr<ResourceDenial> denial(new ResourceDenial());
			denial->code = ResourceDenial::UnknownResource;
			denial->reason = "unknown resource.";
			session->Send(denial);
			return;
		}

		// Check for serializable/cacheable flag.
		if((resource->GetFlags() & (IResource::RFL_CASHEABLE | IResource::RFL_SERIALIZABLE)) == 0)
		{
			// Private resource.
			boost::shared_ptr<ResourceDenial> denial(new ResourceDenial());
			denial->code = ResourceDenial::PrivateResource;
			denial->reason = "private resource.";
			session->Send(denial);
			return;
		}

		// Read the serializable data.
		boost::shared_ptr<ISerializableResource> serial;
		if(resource->GetFlags() & IResource::RFL_SERIALIZABLE)
			serial = boost::dynamic_pointer_cast<ISerializableResource> (resource);

		// Try first with cacheable resources.
		if(resource->GetFlags() & IResource::RFL_CASHEABLE)
		{
			// Ask the cache about the file name.
			boost::shared_ptr<IResourceCache> cache = resourceManager->GetResourceCache();

			std::string path = cache->GetResourcePath(resource);
			if(!path.empty() && IsFile(path))
			{
				// Create the description response.
				boost::shared_ptr<ResourceDescription> description(new ResourceDescription());
				description->resourceId = resource->GetId();
				description->resourceHash = cache->GetResourceHash(resource);
				description->resourceFlags = resource->GetFlags();
				if(serial)
				{
					description->resourceMime = serial->GetMimeType();
					description->resourceExtension = serial->GetFileExtension();
				}
				else
				{
					description->resourceExtension = FileExtension(path);
				}

				// Send the file resource.
				SendFileDescription(description, path);
				return;
			}
		}

		// Now try using serializable resources.
		if(serial && serial->IsMemorySerializable())
		{
			// Create the description response.
			boost::shared_ptr<ResourceDescription> description(new ResourceDescription());
			description->resourceId = resource->GetId();
			description->resourceFlags = resource->GetFlags();
			description->resourceMime = serial->GetMimeType();
			description->resourceExtension = serial->GetFileExtension();

			// Send the memory resource.
			SendMemoryDescription(description, serial);
			return;
		}
	}

	void TransferManager::OnResourceRequestData(boost::shared_ptr<ResourceRequestData> message)
	{
		// Find the resource.
		boost::shared_ptr<ResourceManager> resourceManager = Engine::Get().GetResourceManager();
		boost::shared_ptr<IResource> resource = resourceManager->FindResource(message->resourceId);
		if(!resource)
		{
			// Couldn't find the resource, send a denial message.
			boost::shared_ptr<ResourceDenial> denial(new ResourceDenial());
			denial->code = ResourceDenial::UnknownResource;
			denial->reason = "unknown resource.";
			session->Send(denial);
			return;
		}

		// Check for serializable/cacheable flag.
		if((resource->GetFlags() & (IResource::RFL_CASHEABLE | IResource::RFL_SERIALIZABLE)) == 0)
		{
			// Private resource.
			boost::shared_ptr<ResourceDenial> denial(new ResourceDenial());
			denial->code = ResourceDenial::PrivateResource;
			denial->reason = "private resource.";
			session->Send(denial);
			return;
		}

		// Read the serializable data.
		boost::shared_ptr<ISerializableResource> serial;
		if(resource->GetFlags() & IResource::RFL_SERIALIZABLE)
			serial = boost::dynamic_pointer_cast<ISerializableResource> (resource);

		// Try first with cacheable resources.
		if(resource->GetFlags() & IResource::RFL_CASHEABLE)
		{
			// Ask the cache about the file name.
			boost::shared_ptr<IResourceCache> cache = resourceManager->GetResourceCache();

			std::string path = cache->GetResourcePath(resource);
			if(!path.empty() && IsFile(path))
			{
				// Send the file resource.
				SendFileResource(message->resourceId, path);
				return;
			}
		}

		// Now try using serializable resources.
		if(serial && serial->IsMemorySerializable())
		{
			// Send the memory resource.
			SendMemoryResource(message->resourceId, serial);
			return;
		}
	}

	void TransferManager::SendFileDescription(boost::shared_ptr<ResourceDescription> description, const std::string &path)
	{
		// Open the resource file.
		FILE *file = fopen(path.c_str(), "rb");

		// Get the file size.
		fseek(file, 0, SEEK_END);
		size_t size = ftell(file);
		fseek(file, 0, SEEK_SET);

		// Close the file.
		fclose(file);

		// Calculate the blocks.
		size_t numblocks = size / ResourceData::MaxBlockSize;
		size_t lastBlockSize = size % ResourceData::MaxBlockSize;
		if(lastBlockSize > 0)
			numblocks++;

		// Finish the description.
		description->numblocks = numblocks;
		description->resourceSize = size;
		description->numblocks = numblocks;
		description->blocksize = numblocks > 1 ? ResourceData::MaxBlockSize : lastBlockSize;

		// Send the description.
		session->Send(description);
	}

	void TransferManager::SendMemoryDescription(boost::shared_ptr<ResourceDescription> description, boost::shared_ptr<ISerializableResource> resource)
	{
		// Get the serialized data.
		boost::shared_ptr<IDataFilter> serialData =
				resource->GetSerializedData();

		// Calculate the blocks.
		serialData->Seek(0, IDataFilter::ST_END);
		size_t size = serialData->GetPosition();
		size_t numblocks = size / ResourceData::MaxBlockSize;
		size_t lastBlockSize = size % ResourceData::MaxBlockSize;
		if(lastBlockSize > 0)
			numblocks++;

		// Finish the description.
		description->numblocks = numblocks;
		description->resourceSize = size;
		description->numblocks = numblocks;
		description->blocksize = numblocks > 1 ? ResourceData::MaxBlockSize : lastBlockSize;

		// Send the description.
		session->Send(description);
	}

	void TransferManager::SendFileResource(const uuid &id, const std::string &path)
	{
		// Open the resource file.
		FILE *file = fopen(path.c_str(), "rb");

		// Get the file size.
		fseek(file, 0, SEEK_END);
		size_t size = ftell(file);
		fseek(file, 0, SEEK_SET);

		// Calculate the blocks.
		size_t numblocks = size / ResourceData::MaxBlockSize;
		size_t lastBlockSize = size % ResourceData::MaxBlockSize;
		if(lastBlockSize > 0)
			numblocks++;

		// TODO: implement a callback based transfer.

		// Sending all of the file data.
		boost::shared_ptr<ResourceData> data;
		size_t next = ResourceData::MaxBlockSize;
		for(size_t i = 0; i < numblocks; i++)
		{
			// Create the message.
			data.reset(new ResourceData());

			// Calculate message block size.
			if(i + 1 == numblocks)
				next = lastBlockSize;

			// Fill the message.
			data->resourceId = id;
			data->blockid = i;
			data->blocksize = next;

			// Fill the message payload.
			if(fread(data->block.c_array(), next, 1, file) != 1)
			{
				LogConsole->Warning("Failed to read resource file block.");
				break;
			}

			// Send the message
			session->Send(data);
		}

		// Close the file.
		fclose(file);
	}

	void TransferManager::SendMemoryResource(const uuid &id, boost::shared_ptr<ISerializableResource> resource)
	{
		// Get the serialized data.
		boost::shared_ptr<IDataFilter> serialData =
				resource->GetSerializedData();

		// Calculate the blocks.
		serialData->Seek(0, IDataFilter::ST_END);
		size_t size = serialData->GetPosition();
		size_t numblocks = size / ResourceData::MaxBlockSize;
		size_t lastBlockSize = size % ResourceData::MaxBlockSize;
		if(lastBlockSize > 0)
			numblocks++;

		// TODO: implement a callback based transfer.

		// Send all of the file data.
		serialData->Seek(0, IDataFilter::ST_SET);
		boost::shared_ptr<ResourceData> data;
		size_t next = ResourceData::MaxBlockSize;
		for(size_t i = 0; i < numblocks; i++)
		{
			// Create the message.
			data.reset(new ResourceData());

			// Calculate message block size.
			if(i + 1 == numblocks)
				next = lastBlockSize;

			// Fill the message.
			data->resourceId = id;
			data->blockid = i;
			data->blocksize = next;

			// Fill the message payload.
			serialData->Read(next, data->block.c_array());

			// Send the message
			session->Send(data);
		}
	}

	void TransferManager::OnResourceData(boost::shared_ptr<ResourceData> message)
	{
		// Find the transfer state.
		TransferStates::iterator it = transferStates.find(message->resourceId);
		if(it == transferStates.end())
			return;
		boost::shared_ptr<TransferState> state = it->second;

		// Sanity check.
		assert(state->state == TransferState::ST_RECEIVING);

		// Check the block size.
		if(message->blocksize != state->blockSize &&
			message->blockid + 1 != state->numblocks)
		{
			// Invalid transfer.
			LogConsole->Error("Received invalid resource transfer: block size is wrong.");
			transferStates.erase(it);
			return;
		}
		else if(message->blocksize != state->fileSize % state->blockSize &&
				message->blockid + 1 == state->numblocks &&
				message->blockid > 0)
		{
			// Invalid last block.
			LogConsole->Error("Received invalid resource transfer: last block size is wrong.");
			transferStates.erase(it);
			return;
		}
		else if(message->blocksize != state->blockSize &&
				state->numblocks == 1)
		{
			// Invalid small resource data.
			LogConsole->Error("Received invalid resource transfer: small resource data is wrong.");
			transferStates.erase(it);
			return;
		}

		// Read the data payload.
		size_t position = message->blockid*state->blockSize;
		//state->stream->Seek(position, IDataFilter::ST_SET);
		state->stream->Write(message->blocksize, message->block.c_array());

		// Decrease the remaining blocks.
		state->remainingBlocks--;

		// Check if the download has been finished.
		if(state->remainingBlocks == 0)
		{
			boost::shared_ptr<ResourceManager> resourceManager = Engine::Get().GetResourceManager();
			boost::shared_ptr<IResource> resource;

			// Try with file based loading.
			if(!state->path.empty())
			{
				// Flush the filter.
				state->stream->Flush();

				// Register the file with the cache.
				boost::shared_ptr<IResourceCache> cache = resourceManager->GetResourceCache();
				cache->RegisterResourceFile(state->path);

				try
				{
					// Load the resource.
					resource = resourceManager->LoadFromCache(state->hash);
				}
				catch(const XernException &e)
				{
					LogConsole->Error("Failed to load downloaded resource: %s", e.what());
					return;
				}
			}

			// Now, try with memory based loading.
			if(!resource)
			{
				state->stream->Seek(0, IDataFilter::ST_SET);

				try
				{
					resource = resourceManager->LoadFromMemory(state->mime, state->stream);
				}
				catch(const XernException &e)
				{
					LogConsole->Error("Failed to read downloaded resource: %s", e.what());
					return;
				}
			}

			// Resource loading failed.
			if(!resource)
			{
				LogConsole->Error("Failed to read downloaded resource %s.", boost::lexical_cast<std::string> (message->resourceId).c_str());
				return;
			}

			// Rename the resource.
			resourceManager->RenameResource(resource, message->resourceId);

			// Perform the callbacks.
			TransferState::Callbacks::iterator cit = state->callbacks.begin();
			for(; cit != state->callbacks.end(); cit++)
			{
				boost::shared_ptr<IResourceLoadedCallback> callback = *cit;
				callback->OnLoadedResource(resource);
			}

			// Delete the state.
			transferStates.erase(it);
		}
	}

	void TransferManager::OnResourceDenial(boost::shared_ptr<ResourceDenial> message)
	{
		// Find the transfer state.
		TransferStates::iterator it = transferStates.find(message->resourceId);
		if(it == transferStates.end())
			return;

		// Remove the transfer state.
		transferStates.erase(it);

		// Print the message.
		LogConsole->Error("Failed to download resource[%d]: %s", message->code, message->reason.c_str());
	}

	void TransferManager::OnResourceDescription(boost::shared_ptr<ResourceDescription> message)
	{
		// Find the transfer state.
		TransferStates::iterator it = transferStates.find(message->resourceId);
		if(it == transferStates.end())
			return;

		boost::shared_ptr<TransferState> state = it->second;

		// Try to load the resource from the cache.
		boost::shared_ptr<ResourceManager> resourceMan = Engine::Get().GetResourceManager();
		boost::shared_ptr<IResource> resource;

		try
		{
			resource = resourceMan->LoadFromCache(message->resourceHash);
		}
		catch(const XernException &e)
		{
			LogConsole->Error("Failed to load downloaded resource: %s", e.what());
			return;
		}

		if(resource)
		{
			// Rename the resource.
			resourceMan->RenameResource(resource, message->resourceId);

			try
			{
				// Perform the callbacks.
				TransferState::Callbacks::iterator cit = state->callbacks.begin();
				for(; cit != state->callbacks.end(); cit++)
				{
					boost::shared_ptr<IResourceLoadedCallback> callback = *cit;
					callback->OnLoadedResource(resource);
				}
			}
			catch(const XernException &e)
			{
				LogConsole->Error("Failed to process downloaded resource: %s", e.what());
			}

			// Delete the state.
			transferStates.erase(it);
			return;
		}

		// Read the description.
		state->fileSize = message->resourceSize;
		state->numblocks = message->numblocks;
		state->blockSize = message->blocksize;
		state->remainingBlocks = state->numblocks;
		state->mime = message->resourceMime;

		// Create the stream.
		boost::shared_ptr<IDataFilter> filter;

		// If the resource has an extension, use the cache.
		if(!message->resourceExtension.empty())
		{
			if(message->resourceExtension[0] != '.')
				message->resourceExtension = "." + message->resourceExtension;

			boost::shared_ptr<IResourceCache> cache =
					Engine::Get().GetResourceManager()->GetResourceCache();
			state->path = cache->CreateFileName() + message->resourceExtension;

			// Open the file.
			FILE *file = fopen(state->path.c_str(), "w+b");
			if(!file)
			{
				LogConsole->Error("Failed to open cache file %s", state->path.c_str());
				transferStates.erase(it);
				return;
			}

			// Create the file file.
			filter.reset(new FileFilter(file, true));
		}

		// Create a memory filter.
		if(!filter)
		{
			filter.reset(new BufferFilter(message->resourceSize));
		}

		// Store the filter and start receiving.
		state->stream = filter;
		state->state = TransferState::ST_RECEIVING;

		// Request the resource data.
		boost::shared_ptr<ResourceRequestData> requestData(new ResourceRequestData());
		requestData->resourceId = message->resourceId;
		session->Send(requestData);
	}

};
